import idaapi
import idautils
import idc
import shelve
import time
import json
import os

with open("D:\HGMSim\HGMSimBackEndByFlask\config.json", "r") as file:
    config = json.load(file)


idaLogBasePath = config["featureExtract"]["idaLogBasePath"]
if not os.path.exists(idaLogBasePath):
    os.makedirs(idaLogBasePath)

extract_error_log = os.path.join(
    idaLogBasePath, config["featureExtract"]["idaErrorLog"]
)
if not os.path.isfile(extract_error_log):
    with open(extract_error_log, "w") as f:
        pass

graphInfLog = os.path.join(idaLogBasePath, config["featureExtract"]["idaGraphInfLog"])
if not os.path.isfile(graphInfLog):
    with open(graphInfLog, "w") as f:
        pass

sample_extrace = config["featureExtract"]["idaExtracAllData"]

dataSavePath = config["featureExtract"]["idaFeatureSaveDirectory"]
if not os.path.exists(dataSavePath):
    os.makedirs(dataSavePath)

arithmetric = [
    "add",
    "adc",
    "inc",
    "daa",
    "sub",
    "sbb",
    "dec",
    "cmp",
    "das",
    "mul",
    "imul",
    "div",
    "idiv",
]
dataTrans = ["mov", "push", "pop", "xchg", "xlat", "cwd", "lea", "lds", "les"]
controlTrans = [
    "call",
    "ret",
    "retn",
    "jmp",
    "loop",
    "loope",
    "loopne",
    "int",
    "into",
    "iret",
    "jz",
    "je",
    "jnz",
    "jne",
    "jo",
    "jno",
    "jc",
    "jnc",
    "js",
    "jns",
    "jb",
    "jnae",
    "jnb",
    "jae",
    "ja",
    "jnbe",
    "jna",
    "jbe",
    "jl",
    "jnge",
    "jnl",
    "jge",
    "jg",
    "jnle",
    "jng",
    "jle",
]


def function_call_extract(func, callees, callers):
    """
    保存函数间的调用关系
    """
    func_name = idc.get_func_name(func)
    print("{}被调用".format(func_name))

    # 同一函数可能会多次调用另一个函数
    for ref_ea in idautils.CodeRefsTo(func, 0):
        caller_name = idc.get_func_name(ref_ea)
        print("    {}--{}".format(hex(ref_ea), idc.GetDisasm(ref_ea)))
        print("    {}调用".format(caller_name))
        # 移除未识别出函数的情况
        if caller_name != "":
            callees[caller_name].append(func_name)
            callers[func_name].append(caller_name)


def function_extract(func, callees, callers):
    """
    #增加提取函数属性
    #增加保存函数汇编指令
    """
    func_name = idc.get_func_name(func)
    temp = idaapi.get_func(func)
    start = temp.start_ea
    end = temp.end_ea
    cur_addr = start
    discode = []
    emb = [0] * 8
    # 按照Gemini的方式提取，中心性感觉复杂度很高，目前没有实现
    while cur_addr < end:
        # 保存汇编指令
        discode.append("{}--{}".format(hex(cur_addr), idc.GetDisasm(cur_addr)))
        # 函数内特征
        # 指令数量
        emb[4] += 1

        # 字符串常量,提取存在问题,提取不出来
        insn = idaapi.insn_t()
        if idaapi.decode_insn(insn, cur_addr):
            for operand in insn.ops:
                # 判断操作数是否为立即数操作数
                if operand.type == idaapi.o_imm:
                    # 获取立即数的值
                    immediate_value = operand.value
                    if idc.is_strlit(immediate_value):
                        emb[0] += 1
                        print("字符常量{}".format(idc.GetDisasm(cur_addr)))
                if operand.type == idaapi.o_mem:
                    mem_address = operand.addr
                    print("常量{}".format(idc.GetDisasm(cur_addr)))
                    if idc.is_strlit(mem_address):
                        emb[0] += 1
                        print("字符常量{}".format(idc.GetDisasm(cur_addr)))

        # 数字常量  && 字符串常来给你
        insn = idaapi.insn_t()
        if idaapi.decode_insn(insn, cur_addr) > 0:
            for op in insn.ops:
                if op.type == idaapi.o_imm:
                    emb[1] += 1
                    # print("数值常量{}".format(idc.GetDisasm(cur_addr)))

        # 转移指令
        if idc.print_insn_mnem(cur_addr) in dataTrans:
            emb[2] += 1
        # 控制指令
        elif idc.print_insn_mnem(cur_addr) in controlTrans:
            emb[3] += 1
        # 算术指令
        elif idc.print_insn_mnem(cur_addr) in arithmetric:
            emb[5] += 1
        # 下一条指令起点
        cur_addr = idc.next_head(cur_addr, end)
    # 函数间特征
    # 后代数量
    emb[6] = len(callers[func_name])
    # 上一辈数量
    if func_name in callees:
        emb[7] = len(callees[func_name])
    return emb, discode


def cg_extract(funcs_id, callees, func_num):
    cg_adjmat = []
    for i in range(func_num):
        cg_adjmat.append([0] * func_num)
    for key in callees.keys():
        for calling in callees[key]:
            # 有边则加一
            # cg_adjmat[funcs_id[key]][funcs_id[calling]]+=1
            cg_adjmat[funcs_id[calling]][funcs_id[key]] += 1
    return cg_adjmat


def BB_extract(func):
    cnt = 0
    f = idaapi.FlowChart(idaapi.get_func(func))
    cfg_adjmat = []
    cfg_attr = []
    cfg_dis = []

    for block in f:
        cfg_row = [0] * f.size
        succNum = 0
        predNum = 0
        for succ_block in block.succs():
            cfg_row[succ_block.id] = 1
            succNum += 1
        for i in block.preds():
            predNum += 1

        cur_addr = block.start_ea
        emb = [0] * 8
        emb[6] = succNum
        emb[7] = predNum
        temp = []
        while cur_addr < block.end_ea:  # 遍历汇编代码
            temp.append("{}--{}".format(hex(cur_addr), idc.GetDisasm(cur_addr)))
            if idaapi.is_strlit(idaapi.get_flags(cur_addr)):
                emb[0] += 1
            emb[4] += 1
            if (
                idc.get_operand_type(cur_addr, 0) == 5
                or idc.get_operand_type(cur_addr, 1) == 5
                or idc.get_operand_type(cur_addr, 2) == 5
            ):
                emb[1] += 1
            if idc.print_insn_mnem(cur_addr) in dataTrans:
                emb[2] += 1
            elif idc.print_insn_mnem(cur_addr) in controlTrans:
                emb[3] += 1
            elif idc.print_insn_mnem(cur_addr) in arithmetric:
                emb[5] += 1
            cur_addr = idc.next_head(cur_addr, block.end_ea)
        cfg_dis.append(temp)
        cfg_attr.append(emb)
        cfg_adjmat.append(cfg_row)
    return cfg_adjmat, cfg_attr, cfg_dis


def get_imports_name():
    imports = {}
    current = ""

    def callback(ea, name, oridinal):
        imports[current].append((hex(ea), name, oridinal))
        return True

    muduleNum = idaapi.get_import_module_qty()
    for i in range(0, muduleNum):
        module_name = idaapi.get_import_module_name(i)
        print("模块名称:", module_name)
        current = module_name
        imports[module_name] = []
        idaapi.enum_import_names(i, callback)
    return imports


def controller():
    funcs_id = dict()  # 函数--》id
    callees = dict()  # 被调用者
    callers = dict()  # 被调用者
    funcType = dict()
    func_num = 0
    func_id = 0

    funcs = idautils.Functions()  # 遍历函数的指针
    funcs1 = idautils.Functions()
    funcs_iterator = idautils.Functions()

    # 获得函数及其对应的下标（id）
    for i in funcs_iterator:
        func_name = idc.get_func_name(i)
        flag = "local"
        # 判断是否为静态链接库代码,
        # ida中为蓝色标注为静态连接库函数，粉色的一般为调用动态导入函数的简单thunk函数
        flags = idc.get_func_flags(i)
        if flags & FUNC_LIB:
            flag = "static link"
        funcType.update({func_id: flag})
        funcs_id.update({func_name: func_id})
        func_num += 1
        func_id += 1
        callers[func_name] = []
        callees[func_name] = []

    # 动态导入函数一般不在Functions中，要获取导入表中各个api的名称才行
    apiName = get_imports_name()
    flag = "dynamic import"
    # 加入动态导入函数
    for moduleName in apiName:
        for addr, tarname, temp in apiName[moduleName]:
            addr = idc.get_name_ea_simple(tarname)  # 二进制的api地址
            n_id = func_id
            # Function中同名的Thunk函数，修改函数类别即可
            if tarname in funcs_id:
                n_id = funcs_id[tarname]
                funcType.update({n_id: flag})
            # 不存在api函数，添加！
            else:
                funcType.update({func_id: flag})
                funcs_id.update({tarname: func_id})
                func_num += 1
                func_id += 1
                callers[tarname] = []
                callees[tarname] = []

    # 分别用来保存cfg图，对应的属性，cg的属性和每个函数的汇编代码
    cfg, attr, cattr, discodes, cdis = [], [], [], [], []
    for i in range(len(funcs_id)):
        cfg.append([])
        attr.append([])
        cattr.append([])
        discodes.append([])
        cdis.append([])

    # 构建函数间的调用关系
    for f in funcs:
        func_name = idc.get_func_name(f)
        # 移除动态导入函数的thunk函数，避免重复计算
        if funcType[funcs_id[func_name]] != "dynamic import":
            # callers 保存调用key的所有函数，callees 保存key调用的所有函数
            function_call_extract(f, callees, callers)

    # api的函数调用也需要添加进来
    for moduleName in apiName:
        for addr, tarname, temp in apiName[moduleName]:
            addr = idc.get_name_ea_simple(tarname)
            print(tarname)
            for i in idautils.CodeRefsTo(addr, 0):
                caller_name = idc.get_func_name(i)
                # 移除未识别出函数的情况
                if caller_name != "":
                    print("    {}--{}".format(hex(i), idc.GetDisasm(i)))
                    print("    {}调用".format(caller_name))
                    callees[caller_name].append(tarname)
                    callers[tarname].append(caller_name)

    for f in funcs1:
        func_name = idc.get_func_name(f)
        # 提取函数特征向量
        # 只需要提取functions里面的函数，动态连接的函数无代码实体
        if funcType[funcs_id[func_name]] != "dynamic import":
            emb, discode = function_extract(f, callees, callers)
            print("{}\t--{}\n".format(func_name, emb))
            attr[funcs_id[func_name]] = emb

            if sample_extrace == False:
                discodes[funcs_id[func_name]] = discode
                # 函数基本块的邻接矩阵、特征向量和对应的汇编代码
                cfg_adj, cfg_attr, cfg_dis = BB_extract(f)
                cfg[funcs_id[func_name]] = cfg_adj
                cattr[funcs_id[func_name]] = cfg_attr
                cdis[funcs_id[func_name]] = cfg_dis

    # 获得整个函数的邻接矩阵
    cg = cg_extract(funcs_id, callees, func_num)

    # 保存图的结点信息
    with open(graphInfLog, "a") as f:
        lth = 0
        for temp in cfg:
            lth += len(temp)
        f.write(
            str(idaapi.get_root_filename())
            + "---cfgNum={}  cgNum={}\n".format(lth, len(cg))
        )

    # 保存需要模型训练的信息
    file = shelve.open(dataSavePath + "\\" + str(idaapi.get_root_filename()))
    file["cg"] = cg
    file["cgattr"] = attr
    file["funcType"] = funcType
    file["func_id"] = funcs_id
    if sample_extrace == False:
        file["cfg"] = cfg
        file["cfgattr"] = cattr
        # 应该保存谁调用了它，而不是它调用了谁，聚合邻居
        file["caller"] = callers
        file["callee"] = callees
        file["discodes"] = discodes
        file["cdiscodes"] = cdis
    file.close()
    print("{}数据提取完成".format(time.time()))


def main():
    q = None
    f = None
    try:
        idaapi.auto_wait()
        controller()
        idc.qexit(0)
    except Exception as e:
        with open(extract_error_log, "a+") as file:
            apiName = get_imports_name()
            file.write(apiName + "  " + e + "\n")


if __name__ == "__main__":
    main()
